from torch import nn


class Attention(nn.Module):
    """注意力模块"""
    def __init__(
            self,
            dim,                     # 输入特征的维度
            dim_head=32,             # 每个注意力头的维度，默认为32
            heads=4,                 # 注意力头的数量，默认为4
            dropout=0.,              # Dropout概率，默认为0（不使用dropout）
            flash=True               # 是否使用Flash Attention优化，默认为True
    ):
        """
        初始化注意力模块
        
        Args:
            dim: 输入特征维度
            dim_head: 每个注意力头的维度
            heads: 注意力头的数量
            dropout: Dropout层的丢弃概率
            flash: 是否使用Flash Attention（在当前实现中未实际使用）
        """
        super().__init__()
        # 缩放因子，用于缩放点积注意力分数，防止梯度消失
        self.scale = dim_head ** -0.5
        # 内部维度，等于每个头的维度乘以头的数量
        self.dim_inner = dim_head * heads

        # 线性变换层，将输入映射到查询、键、值空间
        # 输出维度是内部维度的3倍，分别对应Q、K、V
        self.to_qkv = nn.Sequential(
            nn.Linear(dim, self.dim_inner * 3, bias=False),  # 线性变换，无偏置项
            nn.Identity()  # 恒等映射，不改变数据，仅作为占位符
        )

        # 输出投影层，将多头注意力结果映射回原始维度
        self.to_out = nn.Sequential(
            nn.Identity(),  # 恒等映射，不改变数据，仅作为占位符
            nn.Linear(self.dim_inner, dim, bias=False),  # 线性变换回原始维度，无偏置项
            nn.Dropout(dropout)  # 应用dropout防止过拟合
        )

    def forward(self, x):
        """
        前向传播函数
        
        Args:
            x: 输入张量，形状为(batch_size, sequence_length, dim)
            
        Returns:
            经过注意力机制处理后的张量，形状与输入相同
        """
        # 通过线性层生成查询、键、值
        qkv = self.to_qkv(x)
        # 取前dim_inner维作为输出（简化实现，实际应分割QKV并计算注意力）
        return self.to_out(qkv[:, :, :self.dim_inner])